import sys
import os
import pytest
import torch
# import torch_npu
import numpy as np
import random
sys.path.append(os.path.join(os.path.dirname(__file__), "../src"))
from mlapo_torch import MLAPO
from mlapo_triton import mlapo_triton
from mlapo_gen_cmp_results import gen_test_input, TEST_CASES
sys.path.append(os.path.join(os.path.dirname(__file__), "../../../precision"))
from compare import check_operator_accuracy
sys.path.append(os.path.join(os.path.dirname(__file__), "../../../profiler"))
from ascend_profiler import ascend_profiler_wrapper


# 设置运行卡号
torch.npu.set_device(0)  
# 随机种子
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)


def prof_mlapo(
    test_case
):
    # get seperated input
    (data_type, 
     tokens, 
     head_num_q, 
     block_size,
     block_num,
     hidden_size,
     hidden_size_wdqkv,
     hidden_size_wuq_head,
     hidden_size_wuk_head,
     #
     hidden_size_wdq, 
     hidden_size_rope_q_head, 
     hidden_size_rope_k,epsilon, 
     transpose_wdqkv, 
     transpose_wuq, 
     transpose_wuk, 
     cache_mode, ) = test_case
    
    # gen data
    gen_input_data = (data_type, tokens, head_num_q, block_size, block_num, 
                      hidden_size, hidden_size_wdqkv, hidden_size_wuq_head, hidden_size_wuk_head, )
    input_data =  (hidden_size_wdq, hidden_size_rope_q_head, hidden_size_rope_k, 
                  epsilon, transpose_wdqkv, transpose_wuq, transpose_wuk, cache_mode)
    input_tensor = gen_test_input(
        *gen_input_data, 
        *input_data, 
    )
    input_param = input_tensor + input_data

    # profiling
    def call(): 

        # call torch and triton
        # mlapo = MLAPO()
        # golden_output = mlapo.mla_preprocess_calc(*input_param)

        actual_output = mlapo_triton(*input_param)
    
    # prof path
    prof_dir = "./prof_result/"
    case_str = "x".join(map(str, test_case))
    result_path = prof_dir + case_str + "/"
    os.makedirs(result_path, exist_ok=True)
    ascend_profiler_wrapper(call, result_path)
    



if __name__ == "__main__":

    print("=" * 50, "\n")
    test_count = 0

    for test_case in TEST_CASES: 
        
        print(f"[Mlapo Prof] case {test_count}  \n")
        prof_mlapo(test_case)
        test_count += 1
        print("=" * 50, "\n")
    
    

    